// --------------------------------------------------------------------------------------------------------------------
// <copyright file="ObjectDetectorPicknPlaceIntegratorState.cs" company="Microsoft Corporation">
// The MIT License (MIT)
// 
// Copyright (c) 2014, Microsoft Corporation
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//  THE SOFTWARE.
// </copyright>
// --------------------------------------------------------------------------------------------------------------------

namespace Test.Robotics.Manipulation
{
    using System;
    using System.Collections.Generic;
    using System.IO;
    using System.Windows;
    using Microsoft.Robotics.FiniteStateMachine;
    using Microsoft.Robotics.Manipulation.PicknPlace;
    using Microsoft.Robotics.Numerics;
    using Microsoft.Robotics.PhysicalObjects;
    using Microsoft.Robotics.Tracing;
    using Microsoft.Robotics.Vision;
    using Microsoft.Robotics.Vision.ObjectDetection;
    using Microsoft.Robotics.Vision.Runtime;

    /// <summary>
    /// Maintains and executes state updates for the object detector pick and place integrator agent
    /// </summary>
    public class ObjectDetectorPicknPlaceIntegratorState : IDisposable
    {
        /// <summary>
        /// Maximum number of tries before declare failure
        /// </summary>
        private readonly int maxRetry;

        /// <summary>
        /// Minimum number of detection results before checking result
        /// </summary>
        private readonly int minDetectionCount;

        /// <summary>
        /// Count of object detection messages received since last pickMessage
        /// </summary>
        private long pickMessageOriginatingTime;

        /// <summary>
        /// Count of detection check retries
        /// </summary>
        private int retryCount;

        /// <summary>
        /// Master data log file
        /// </summary>
        private string masterLogPathName;

        /// <summary>
        /// Path to kinect data files. For example d:\dataFile\Kinect.txt
        /// The final name is constructed as  d:\dataFile\Kinect_dateStamp.txt
        /// </summary>
        private string kinectFilesPath;

        /// <summary>
        /// Master log file stream
        /// </summary>
        private TextWriter masterLogFileStream;

        /// <summary>
        /// Count of images saved
        /// </summary>
        private int savedImageCount;

        /// <summary>
        /// Counts of pick object s in last pick and move request
        /// </summary>
        private Dictionary<string, int> lastPickObjectCounts;

        /// <summary>
        /// Initializes a new instance of the <see cref="ObjectDetectorPicknPlaceIntegratorState" /> class
        /// </summary>
        /// <param name="minDetectionCount">Minimum number of results before checking result validity</param>
        /// <param name="maxRetry">Maximum number of detection retries</param>
        /// <param name="masterLogPathName">Path to master log file</param>
        /// <param name="kinectFilesPath">Path name for kinect data files</param>
        public ObjectDetectorPicknPlaceIntegratorState(
            int minDetectionCount,
            int maxRetry,
            string masterLogPathName,
            string kinectFilesPath)
        {
            this.minDetectionCount = minDetectionCount;
            this.maxRetry = maxRetry;
            this.LastPickStatus = PicknPlaceStatusEnum.NotStarted;
            this.LastStatusString = string.Empty;
            this.pickMessageOriginatingTime = 0;
            this.savedImageCount = 0;
            this.masterLogPathName = masterLogPathName;
            this.kinectFilesPath = kinectFilesPath;
            this.masterLogFileStream = null;
        }

        /// <summary>
        /// Gets the last status string
        /// </summary>
        public string LastStatusString { get; private set; }

        /// <summary>
        /// Gets the last recorded pick status
        /// </summary>
        public PicknPlaceStatusEnum LastPickStatus { get; private set; }

        /// <summary>
        /// Gets the last received object detection result
        /// </summary>
        public KinectObjectDetectorResultIRMessage LastReceivedObjectDetectionResult { get; private set; }

        /// <summary>
        /// Gets Last pick and move requests received
        /// </summary>
        public PickAndMove[] LastPickRequests { get; private set; }

        /// <summary>
        /// Gets the last pick mode
        /// </summary>
        public PicknPlaceModeEnum LastPickMode { get; private set; }

        /// <summary>
        /// Gets a value indicating whether we are ready to check object detection
        /// </summary>
        public bool IsDetectionReady
        {
            get
            {
                return ((this.LastPickStatus != PicknPlaceStatusEnum.PlaceInProgress) &&
                        (this.LastPickStatus != PicknPlaceStatusEnum.PlaceComplete) &&
                        (this.LastPickStatus != PicknPlaceStatusEnum.PickFailFalseNegative) &&
                        (this.LastPickStatus != PicknPlaceStatusEnum.PickFailFalsePositive) &&
                        (this.LastPickStatus != PicknPlaceStatusEnum.PickFailLargeError) &&
                        (this.pickMessageOriginatingTime > 0) &&
                        (null != this.LastPickRequests) &&
                        (this.LastPickRequests.Length > 0) &&
                        (this.LastReceivedObjectDetectionResult.OriginatingTime > this.pickMessageOriginatingTime) &&
                        (this.retryCount < this.maxRetry));
            }
        }

        /// <summary>
        /// Gets a value indicating whether visual detection failed for last pick target
        /// </summary>
        public bool IsVisualDetectionFail
        {
            get
            {
                return false == string.IsNullOrEmpty(this.LastStatusString) &&
                    (this.LastPickStatus == PicknPlaceStatusEnum.PickFailLargeError ||
                     this.LastPickStatus == PicknPlaceStatusEnum.PickFailLowConfidence ||
                     this.LastPickStatus == PicknPlaceStatusEnum.PickFailFalsePositive ||
                     this.LastPickStatus == PicknPlaceStatusEnum.PickFailFalseNegative);
            }
        }

        /// <summary>
        /// Gets destination of the move message.
        /// </summary>
        public Pose LastMoveToPose { get; private set; }

        /// <summary>
        /// Extract detected pose estimates for a target object that are above the confidence limit
        /// from a candidates list
        /// </summary>
        /// <param name="poseEstimates">Candidate pose estimates</param>
        /// <param name="objectDescription">Target object description</param>
        /// <param name="confThreshold">Filtering confidence threshold. If less than 0 use the threshold pose threshold</param>
        /// <returns>List of poses meeting the criteria</returns>
        public static List<ObjectPoseEstimate> GetDetectEstimatesForObject(ObjectPoseEstimate[] poseEstimates, string objectDescription, double confThreshold = -1.0)
        {
            List<ObjectPoseEstimate> detectList = new List<ObjectPoseEstimate>();

            foreach (ObjectPoseEstimate est in poseEstimates)
            {
                double threshold = (confThreshold < 0.0) ? est.ConfidenceThreshold : confThreshold;
                if (est.Confidence >= threshold && 0 == string.Compare(est.ManipulatedObjectDescription.Description, objectDescription, true))
                {
                    detectList.Add(est);
                }
            }

            return detectList;
        }

        /// <summary>
        /// Return rotation from a quaternion
        /// </summary>
        /// <param name="q">Input quaternion</param>
        /// <returns>Rotation angle</returns>
        public static double AngleFromQuaternion(Quaternion q)
        {
            AxisAngle a = Quaternion.ToAxisAngle(Quaternion.Normalize(q));
            double angle = a.Angle * MathConstants.Radians2Degrees;
            return angle;
        }

        /// <summary>
        /// Update when a new pick message is received. Resets accumulators
        /// and creates a list of object to be picked
        /// </summary>
        /// <param name="pickMoveRequests">The pick and move requests in the move</param>
        /// <param name="messageTime">Originating time for pick message</param>
        /// <param name="expectedObjects">List of expected objects in ground truth</param>
        public void NewPickMessageReceived(PickAndMove[] pickMoveRequests, long messageTime, List<string> expectedObjects)
        {
            if (true == this.CheckPickMoveRequest(pickMoveRequests, expectedObjects))
            {
                this.LastPickRequests = pickMoveRequests;
                this.pickMessageOriginatingTime = messageTime;
                this.LastPickStatus = PicknPlaceStatusEnum.PickRequestReceived;
                this.retryCount = 0;
                this.LastStatusString = string.Empty;
                this.lastPickObjectCounts = new Dictionary<string, int>();

                foreach (PickAndMove request in pickMoveRequests)
                {
                    string objectDesc = request.ManipulatedObjectDescription.Description;
                    if (true == this.lastPickObjectCounts.ContainsKey(objectDesc))
                    {
                        this.lastPickObjectCounts[objectDesc]++;
                    }
                    else
                    {
                        this.lastPickObjectCounts.Add(objectDesc, 1);
                    }

                    this.LastPickMode = request.PickMode;
                }
            }
            else
            {
                this.LastPickRequests = null;
                this.LastPickStatus = PicknPlaceStatusEnum.PickFailFalseNegative;
            }
        }

        /// <summary>
        /// Update status when a new place message is received
        /// </summary>
        /// <param name="placePose">Requested object pose in place message</param>
        /// <param name="moveToPose">Requested move to pose</param>
        /// <param name="manipulatedObjectDescription">Requested object description</param>
        public void NewPlaceMessageReceived(Pose placePose, Pose moveToPose, PhysicalObjectDescription manipulatedObjectDescription)
        {
            this.LastPickStatus = PicknPlaceStatusEnum.PlaceRequestReceived;
        }

        /// <summary>
        /// Handles reception of new object detection result by updating  detection status
        /// and checking the detection against the pick request.
        /// Note currently success is partially driven by knowing the true object position
        /// </summary>
        /// <param name="message">Message containing the detection result</param>
        /// <param name="pickAndMoves">List of matching pick and move requests</param>
        /// <returns>True if the new detection result is successful, false otherwise. Success means the detection result contains
        /// all requested pick objects and the detected pose are sufficiently close to requested pick locations</returns>
        public bool IsNewDectionSuccess(KinectObjectDetectorResultIRMessage message, out PickAndMove[] pickAndMoves)
        {
            bool isSuccess = false;
            this.LastReceivedObjectDetectionResult = message;
            pickAndMoves = null;

            if (true == this.IsDetectionReady && message.State.Length > 0)
            {
                pickAndMoves = new PickAndMove[this.LastPickRequests.Length];
                this.retryCount++;
                isSuccess = true;

                for (int i = 0; i < this.LastPickRequests.Length; ++i)
                {
                    bool itemSuccess = false;
                    PickAndMove pickRequest = this.LastPickRequests[i];
                    string requestObject = pickRequest.ManipulatedObjectDescription.Description;
                    List<ObjectPoseEstimate> estimates = GetDetectEstimatesForObject(message.State, requestObject);
                    int expectCount = this.lastPickObjectCounts[requestObject];

                    if (estimates.Count != expectCount)
                    {
                        itemSuccess = false;
                        this.LastStatusString = string.Format("Expect {0} objects of type {1} Found {2}", expectCount, requestObject, estimates.Count);
                        TraceOut.Warning(TraceContexts.Manip, this.LastStatusString);

                        if (this.retryCount >= this.maxRetry)
                        {
                            this.LastPickStatus = (estimates.Count > expectCount) ? PicknPlaceStatusEnum.PickFailFalsePositive : PicknPlaceStatusEnum.PickFailFalseNegative;
                        }
                    }
                    else
                    {
                        Vector3 target = pickRequest.PickFromPose.Position;
                        estimates.Sort((a, b) => { return (a.GlobalPose.Position - target).Length().CompareTo((b.GlobalPose.Position - target).Length()); });
                        itemSuccess = this.IsThisDetectSuccess(estimates[0], pickRequest.PickFromPose);

                        ObjectPoseEstimate bestEstimate = estimates[0];
                        pickAndMoves[i] = new PickAndMove(
                            bestEstimate.ManipulatedObjectDescription,
                            bestEstimate.GlobalPose,
                            pickRequest.MoveToPose,
                            pickRequest.PickMode);
                    }

                    isSuccess &= itemSuccess;
                }

                this.WriteImages(message);
                if (true == isSuccess)
                {
                    // Done with set of requests
                    this.pickMessageOriginatingTime = 0;
                    this.LastStatusString = null;
                    this.LastPickStatus = PicknPlaceStatusEnum.PickInProgress;
                }
            }

            return isSuccess;
        }

        /// <summary>
        /// Implement IDisposable
        /// </summary>
        public void Dispose()
        {
            this.Dispose(true);
            GC.SuppressFinalize(this);
        }

        /// <summary>
        /// Dispose resources
        /// </summary>
        /// <param name="disposing">True if dispose is called by user code, false otherwise</param>
        protected virtual void Dispose(bool disposing)
        {
            if (disposing)
            {
                if (null != this.masterLogFileStream)
                {
                    this.masterLogFileStream.Dispose();
                    this.masterLogFileStream = null;
                }
            }
        }

        /// <summary>
        /// Check if a single object candidate pose matches the requested pose
        /// </summary>
        /// <param name="objectPose">Detected candidate pose</param>
        /// <param name="requestPos">Requested target pose</param>
        /// <returns>True if candidate estimate is good estimate of target</returns>
        private bool IsThisDetectSuccess(ObjectPoseEstimate objectPose, Pose requestPos)
        {
            bool isSuccess = false;

            Vector3 estimatePos = objectPose.GlobalPose.Position;
            Vector3 actualPos = requestPos.Position;
            double estimateRot = AngleFromQuaternion(objectPose.GlobalPose.Orientation);
            double actualRot = AngleFromQuaternion(requestPos.Orientation);

            // Change
            Vector3 diffPos = estimatePos - actualPos;
            double diffRot = Math.Abs(estimateRot - actualRot);
            diffRot = Math.Min(diffRot, Math.Abs(diffRot - 180.0));

            TraceOut.Info(TraceContexts.Manip, "Detected {0} Conf. {1:F4}\tEstimPos {2:F4} {3:F4} {4:F4} {5:F4} ", objectPose.ManipulatedObjectDescription.Description, objectPose.Confidence, estimatePos.X, estimatePos.Y, estimatePos.Z, estimateRot);
            TraceOut.Info(TraceContexts.Manip, "\t ActualPos {0:F4} {1:F4} {2:F4} {3:F4}", actualPos.X, actualPos.Y, actualPos.Z, actualRot);
            TraceOut.Info(TraceContexts.Manip, "\tDiffPos {0:F4} {1:F4} {2:F4} {3:F4}", diffPos.X, diffPos.Y, diffPos.Z, diffRot);

            // Check for success
            if ((this.retryCount <= this.maxRetry) &&
              (objectPose.Confidence >= objectPose.ConfidenceThreshold) &&
              ((Math.Abs(diffPos.X) < 0.01) && (Math.Abs(diffPos.Y) < 0.013) && (diffPos.Z > -0.009 && diffPos.Z < 0.01)) &&
                (Math.Abs(diffRot) < 15))
            {
                isSuccess = true;
            }
            else if (this.retryCount >= this.maxRetry)
            {
                if (objectPose.Confidence < objectPose.ConfidenceThreshold)
                {
                    this.LastStatusString = string.Format(
                        "Detection confidence = {0:F4} below threshold {1:F4} Placement Diff {{ {2:F4} {3:F4} {4:F4} }}",
                        objectPose.Confidence,
                        objectPose.ConfidenceThreshold,
                        diffPos.X,
                        diffPos.Y,
                        diffPos.Z);
                    this.LastPickStatus = PicknPlaceStatusEnum.PickFailLowConfidence;
                    TraceOut.Info(TraceContexts.Manip, "Detection Failed {0}", this.LastStatusString);
                }
                else
                {
                    this.LastStatusString = string.Format("Visual detection error too high {{ {0:F4} {1:F4} {2:F4} }} Conf {3:F4}", diffPos.X, diffPos.Y, diffPos.Z, objectPose.Confidence);
                    this.LastPickStatus = PicknPlaceStatusEnum.PickFailLargeError;
                    TraceOut.Info(TraceContexts.Manip, "Detection Failed {0}", this.LastStatusString);
                }
            }
            else
            {
                TraceOut.Info(TraceContexts.Manip, "No decision yet number of attempts {0}", this.retryCount);
            }

            return isSuccess;
        }

        /// <summary>
        /// Write kinect depth and IR images
        /// <param name="message">Detection result message containing images to save</param>
        /// </summary>
        /// <param name="message">Latest object detection message</param>
        private void WriteImages(KinectObjectDetectorResultIRMessage message)
        {
            if (true == string.IsNullOrEmpty(this.kinectFilesPath) || null == this.masterLogPathName)
            {
                return;
            }

            if (this.savedImageCount <= 0)
            {
                this.InitializeForSaving(message.State[0].ManipulatedObjectDescription);
            }

            string kinectFinalName = Path.GetFileNameWithoutExtension(this.kinectFilesPath) + DateTime.Now.ToString("s").Replace(':', '_') + Path.GetExtension(this.kinectFilesPath);
            kinectFinalName = Path.Combine(Path.GetDirectoryName(this.kinectFilesPath), kinectFinalName);
            if (true == File.Exists(kinectFinalName))
            {
                return;
            }

            using (FileStream fs = new FileStream(kinectFinalName, FileMode.CreateNew))
            {
                this.SaveIRKinectImageBinary(message.IRframe, fs);
                this.SaveDepthImageBinary(message.DepthFrame, fs);
            }

            foreach (PickAndMove pickMove in this.LastPickRequests)
            {
                Pose objectPose = pickMove.PickFromPose;
                PhysicalObjectDescription objDesc = pickMove.ManipulatedObjectDescription;

                Vector3 p = objectPose.Position;
                Quaternion q = objectPose.Orientation;

                this.masterLogFileStream.Write("{0}", kinectFinalName);
                this.masterLogFileStream.Write("\t{0}\t{1}\t{2}\t{3}", objDesc.Description, objDesc.WidthInM, objDesc.HeightInM, objDesc.DepthInM);
                this.masterLogFileStream.Write("\t{0}\t{1}\t{2}", p.X, p.Y, p.Z);
                this.masterLogFileStream.Write("\t{0}\t{1}\t{2}\t{3}", q.X, q.Y, q.Z, q.W);

                TraceOut.Info(TraceContexts.Manip, "saved {0} at {1:F4} {2:F4} {3:F4}", kinectFinalName, p.X, p.Y, p.Z);
                this.masterLogFileStream.WriteLine();
                this.masterLogFileStream.Flush();
                this.savedImageCount++;
            }
        }

        /// <summary>
        /// Save kinect IR image in Binary. These can be used for offline metrics and retraining if necessary
        /// </summary>
        /// <param name="irFrame">IR image frame</param>
        /// <param name="fs">Saving stream</param>
        private void SaveIRKinectImageBinary(ImageFrame<ushort> irFrame, FileStream fs)
        {
            if (null != irFrame)
            {
                BinaryWriter bw = new BinaryWriter(fs);

                ushort[] irData = irFrame.ImageData;
                bw.Write(irFrame.Width);
                bw.Write(irFrame.Height);
                bw.Write(irFrame.BytesPerPixel);
                for (int row = 0; row < irFrame.Height; ++row)
                {
                    for (int col = 0; col < irFrame.Width; ++col)
                    {
                        bw.Write(irFrame[col, row]);
                    }
                }
            }
        }

        /// <summary>
        /// Save the depth image as Binary
        /// </summary>
        /// <param name="depthFrame">Depth image frame</param>
        /// <param name="fs">Saving stream</param>
        private void SaveDepthImageBinary(ImageFrameDepth depthFrame, FileStream fs)
        {
            if (null != depthFrame)
            {
                BinaryWriter bw = new BinaryWriter(fs);

                bw.Write(depthFrame.Width);
                bw.Write(depthFrame.Height);
                bw.Write(depthFrame.BytesPerPixel);
                for (int idx = 0; idx < depthFrame.Width * depthFrame.Height; ++idx)
                {
                    bw.Write(depthFrame[idx]);
                }
            }
        }

        /// <summary>
        /// Checks if the move request contains all expected objects and no extras
        /// </summary>
        /// <param name="pickMoveRequests">Pick and move request</param>
        /// <param name="expectedObjects">Set of expected object descriptions</param>
        /// <returns>True if all objects are found</returns>
        private bool CheckPickMoveRequest(PickAndMove[] pickMoveRequests, List<string> expectedObjects)
        {
            if (null == expectedObjects)
            {
                return true;
            }

            bool didPass = false;

            if (pickMoveRequests.Length == expectedObjects.Count)
            {
                int matchCount = 0;

                foreach (string expect in expectedObjects)
                {
                    PickAndMove found = Array.Find(pickMoveRequests, p => { return string.Compare(p.ManipulatedObjectDescription.Description, expect, true) == 0; });
                    if (null == found)
                    {
                        TraceOut.Info(TraceContexts.Manip, "Did not find expected object {0}", expect);
                        break;
                    }

                    ++matchCount;
                }

                didPass = (matchCount == expectedObjects.Count);
            }
            else
            {
                TraceOut.Info(TraceContexts.Manip, "Missing objects in received Pick request Expect {0} objects found {1}", expectedObjects.Count, pickMoveRequests.Length);
            }

            return didPass;
        }

        /// <summary>
        /// Initialize - set up path names
        /// </summary>
        /// <param name="objectDescription">Object used to set up path names</param>
        private void InitializeForSaving(PhysicalObjectDescription objectDescription)
        {
            string dateStamp = DateTime.Now.ToString("s").Replace(':', '_');
            string objType = "_" + (objectDescription == null ? string.Empty : objectDescription.Description) + "_";
            string masterFile = Path.GetFileNameWithoutExtension(this.masterLogPathName) + objType + dateStamp + Path.GetExtension(this.masterLogPathName);
            string dirName = Path.GetDirectoryName(this.masterLogPathName);
            if (false == Directory.Exists(dirName))
            {
                Directory.CreateDirectory(dirName);
            }

            this.masterLogPathName = Path.Combine(dirName, masterFile);
            this.masterLogFileStream = new StreamWriter(this.masterLogPathName, false);

            dirName = Path.GetDirectoryName(this.kinectFilesPath);
            dirName = dirName + objType + dateStamp;
            this.kinectFilesPath = Path.Combine(dirName, Path.GetFileName(this.kinectFilesPath));
            if (false == Directory.Exists(dirName))
            {
                Directory.CreateDirectory(dirName);
            }
        }
    }
}
